"""
Modbus TCP Connector Implementation

Uses pymodbus library for asynchronous Modbus TCP client operations.
Supports Modbus TCP protocol for industrial automation and process control.
"""

from __future__ import annotations

import asyncio
from datetime import datetime
from typing import Any

from pymodbus.client import AsyncModbusTcpClient
from pymodbus.exceptions import ConnectionException, ModbusException

# Import framer classes used by pymodbus. These are passed as the `framer` argument
# to the AsyncModbusTcpClient constructor (it expects a Framer class/type, not a str).
# pymodbus 3.x renamed framers: ModbusXxxFramer -> FramerXxx
from pymodbus.framer import FramerAscii, FramerRTU, FramerSocket

from app.connectors.base import BaseConnector, ConnectionStatus, DataPoint, DataQuality
from app.connectors.exceptions import (
    ConnectionError as ConnectorConnectionError,
)
from app.connectors.exceptions import (
    ConnectionTimeoutError,
    DeviceOfflineError,
    ReadError,
    WriteError,
)


# Map simple framer names to the actual framer classes expected by pymodbus.
# Keep this at module-level so tests or other modules can reuse it if needed.
FRAMER_MAP: dict[str, type] = {
    "socket": FramerSocket,
    "rtu": FramerRTU,
    "ascii": FramerAscii,
}


class ModbusTCPConnector(BaseConnector):
    """
    Modbus TCP protocol connector

    Supports:
    - Modbus TCP/IP protocol
    - Read/write coils (function codes 01, 05, 15)
    - Read/write discrete inputs (function code 02)
    - Read/write holding registers (function codes 03, 06, 16)
    - Read input registers (function code 04)
    - Multiple slave devices (unit_id)

    Configuration:
        host (str): Modbus TCP server host/IP address
        port (int, optional): Modbus TCP server port (default: 502)
        unit_id (int, optional): Modbus slave unit ID (default: 1)
        timeout (int, optional): Connection timeout in seconds (default: 5)
        framer (str, optional): Modbus frame type ("socket", "rtu", "ascii", default: "socket")
    """

    def __init__(self, config: dict[str, Any]) -> None:
        """
        Initialize Modbus TCP connector

        Args:
            config: Connector configuration
                Required:
                    - device_id: Unique device identifier
                    - host: Modbus TCP server host/IP
                Optional:
                    - port: Server port (default: 502)
                    - unit_id: Slave unit ID (default: 1)
                    - timeout: Connection timeout (default: 5)
                    - framer: Frame type (default: "socket")
        """
        # Validate Modbus specific configuration BEFORE calling super().__init__()
        if "host" not in config:
            raise ValueError("Missing required config key: host")

        # Initialize Modbus specific attributes
        self.host: str = config["host"]
        self.port: int = config.get("port", 502)
        self.unit_id: int = config.get("unit_id", 1)
        self.framer: str = config.get("framer", "socket")

        # Modbus client instance
        self._client: AsyncModbusTcpClient | None = None

        # Call parent constructor
        super().__init__(config)

        self.logger.info(f"Modbus TCP connector initialized for {self.host}:{self.port}")

    def _validate_config(self) -> None:
        """
        Validate Modbus TCP specific configuration

        Raises:
            ValueError: If configuration is invalid
        """
        super()._validate_config()

        # Validate port range
        if not (1 <= self.port <= 65535):
            raise ValueError(f"Invalid port: {self.port}. Must be between 1 and 65535")

        # Validate unit_id
        if not (0 <= self.unit_id <= 255):
            raise ValueError(f"Invalid unit_id: {self.unit_id}. Must be between 0 and 255")

        # Validate framer type
        valid_framers = ["socket", "rtu", "ascii"]
        if self.framer not in valid_framers:
            raise ValueError(
                f"Invalid framer: {self.framer}. Must be one of: {', '.join(valid_framers)}"
            )

    async def connect(self) -> bool:
        """
        Establish connection to Modbus TCP server

        Returns:
            bool: True if connection successful

        Raises:
            ConnectionError: If connection fails
            ConnectionTimeoutError: If connection times out
        """
        try:
            self.status = ConnectionStatus.CONNECTING
            self.logger.info(f"Connecting to Modbus TCP server: {self.host}:{self.port}")

            # Resolve framer string to the actual framer class expected by pymodbus
            framer_class = FRAMER_MAP.get(self.framer)
            if framer_class is None:
                raise ValueError(f"Unsupported framer type: '{self.framer}'")

            # Create client instance - pass the framer class (not a string)
            self._client = AsyncModbusTcpClient(
                host=self.host,
                port=self.port,
                framer=framer_class,
                timeout=self.timeout,
            )

            # Connect to server
            connected = await self._client.connect()

            if not connected:
                raise ConnectorConnectionError(
                    f"Failed to connect to {self.host}:{self.port}"
                )

            # Verify connection by reading a test register (address 0)
            # This also validates that the device is responsive
            try:
                await self._client.read_holding_registers(address=0, count=1, slave=self.unit_id)
            except Exception as e:
                self.logger.warning(f"Initial test read failed: {e}")
                # Don't fail connection if test read fails - device might not have register 0

            self.status = ConnectionStatus.CONNECTED
            self.logger.info(f"Connected to Modbus TCP server at {self.host}:{self.port}")
            return True

        except TimeoutError as e:
            self.status = ConnectionStatus.ERROR
            self.logger.error(f"Connection timeout: {e}")
            raise ConnectionTimeoutError(
                f"Connection to {self.host}:{self.port} timed out after {self.timeout}s"
            ) from e

        except ConnectionException as e:
            self.status = ConnectionStatus.ERROR
            self.logger.error(f"Modbus connection error: {e}")
            raise ConnectorConnectionError(
                f"Failed to connect to {self.host}:{self.port}: {e}"
            ) from e

        except Exception as e:
            self.status = ConnectionStatus.ERROR
            self.logger.error(f"Unexpected error during connection: {e}")
            raise ConnectorConnectionError(
                f"Failed to connect to {self.host}:{self.port}: {e}"
            ) from e

    async def disconnect(self) -> bool:
        """
        Disconnect from Modbus TCP server

        Returns:
            bool: True if disconnection successful
        """
        try:
            if self._client is not None:
                self.logger.info("Disconnecting from Modbus TCP server")
                self._client.close()
                self._client = None

            self.status = ConnectionStatus.DISCONNECTED
            self.logger.info("Disconnected from Modbus TCP server")
            return True

        except Exception as e:
            self.logger.error(f"Error during disconnection: {e}")
            self.status = ConnectionStatus.ERROR
            return False

    async def read_single(self, address: str) -> DataPoint:
        """
        Read a single data point from Modbus TCP server

        Args:
            address: Modbus address in format:
                     "coil:<addr>" - Read coil (function code 01)
                     "discrete:<addr>" - Read discrete input (function code 02)
                     "holding:<addr>" - Read holding register (function code 03)
                     "input:<addr>" - Read input register (function code 04)
                     Default: holding register if no prefix

        Returns:
            DataPoint: Read data point with value, timestamp, and quality

        Raises:
            ReadError: If read operation fails
            DeviceOfflineError: If device is not connected
        """
        if self.status != ConnectionStatus.CONNECTED or self._client is None:
            raise DeviceOfflineError(
                f"Device {self.name} is not connected. Current status: {self.status}"
            )

        try:
            # Parse address format
            register_type, reg_address = self._parse_address(address)

            # Perform read based on register type
            result = await self._read_register(register_type, reg_address)

            # Create DataPoint
            data_point = DataPoint(
                value=result["value"],
                timestamp=datetime.now(),
                quality=DataQuality.GOOD if result["success"] else DataQuality.BAD,
                address=address,
                metadata={
                    "register_type": register_type,
                    "unit_id": self.unit_id,
                },
            )

            self.logger.debug(f"Read from {address}: {data_point.value}")
            return data_point

        except ModbusException as e:
            self.logger.error(f"Modbus error reading {address}: {e}")
            raise ReadError(f"Failed to read from {address}: {e}") from e

        except Exception as e:
            self.logger.error(f"Unexpected error reading {address}: {e}")
            raise ReadError(f"Failed to read from {address}: {e}") from e

    async def read_multiple(self, addresses: list[str]) -> list[DataPoint]:
        """
        Read multiple data points from Modbus TCP server

        Args:
            addresses: List of Modbus addresses

        Returns:
            List[DataPoint]: List of read data points

        Raises:
            ReadError: If read operation fails
        """
        if self.status != ConnectionStatus.CONNECTED or self._client is None:
            raise DeviceOfflineError(
                f"Device {self.name} is not connected. Current status: {self.status}"
            )

        try:
            # Read all addresses in parallel
            read_tasks = [self.read_single(addr) for addr in addresses]
            data_points = await asyncio.gather(*read_tasks, return_exceptions=True)

            # Filter out exceptions and log errors
            valid_data_points: list[DataPoint] = []
            for i, result in enumerate(data_points):
                if isinstance(result, Exception):
                    self.logger.error(f"Failed to read {addresses[i]}: {result}")
                    # Create error data point
                    valid_data_points.append(
                        DataPoint(
                            value=None,
                            timestamp=datetime.now(),
                            quality=DataQuality.BAD,
                            address=addresses[i],
                            metadata={"error": str(result)},
                        )
                    )
                else:
                    valid_data_points.append(result)  # type: ignore

            self.logger.debug(f"Batch read {len(addresses)} addresses")
            return valid_data_points

        except Exception as e:
            self.logger.error(f"Error in batch read: {e}")
            raise ReadError(f"Failed to read multiple addresses: {e}") from e

    async def write_single(self, address: str, value: Any) -> bool:
        """
        Write a single value to Modbus TCP server

        Args:
            address: Modbus address (same format as read_single)
            value: Value to write (int for registers, bool for coils)

        Returns:
            bool: True if write successful

        Raises:
            WriteError: If write operation fails
            DeviceOfflineError: If device is not connected
        """
        if self.status != ConnectionStatus.CONNECTED or self._client is None:
            raise DeviceOfflineError(
                f"Device {self.name} is not connected. Current status: {self.status}"
            )

        try:
            # Parse address format
            register_type, reg_address = self._parse_address(address)

            # Check if register type is read-only
            if register_type in ["discrete", "input"]:
                raise WriteError(f"Cannot write to {register_type} registers (read-only)")

            # Perform write based on register type
            success = await self._write_register(register_type, reg_address, value)

            if not success:
                raise WriteError(f"Write operation returned failure for {address}")

            self.logger.debug(f"Wrote to {address}: {value}")
            return True

        except WriteError:
            # Re-raise WriteError as-is
            raise

        except ModbusException as e:
            self.logger.error(f"Modbus error writing to {address}: {e}")
            raise WriteError(f"Failed to write to {address}: {e}") from e

        except Exception as e:
            self.logger.error(f"Unexpected error writing to {address}: {e}")
            raise WriteError(f"Failed to write to {address}: {e}") from e

    async def write_multiple(self, data: dict[str, Any]) -> bool:
        """
        Write multiple values to Modbus TCP server

        Args:
            data: Dictionary of {address: value} pairs

        Returns:
            bool: True if all writes successful

        Raises:
            WriteError: If any write operation fails
        """
        if self.status != ConnectionStatus.CONNECTED or self._client is None:
            raise DeviceOfflineError(
                f"Device {self.name} is not connected. Current status: {self.status}"
            )

        try:
            # Write all addresses in parallel
            write_tasks = [self.write_single(addr, value) for addr, value in data.items()]
            results = await asyncio.gather(*write_tasks, return_exceptions=True)

            # Check for failures
            failed_writes = [
                addr for addr, result in zip(data.keys(), results)
                if isinstance(result, Exception)
            ]

            if failed_writes:
                raise WriteError(f"Failed to write to addresses: {', '.join(failed_writes)}")

            self.logger.debug(f"Batch wrote {len(data)} addresses")
            return True

        except Exception as e:
            self.logger.error(f"Error in batch write: {e}")
            raise WriteError(f"Failed to write multiple addresses: {e}") from e

    async def health_check(self) -> bool:
        """
        Perform health check by reading a test register

        Returns:
            bool: True if server is reachable and healthy
        """
        if self.status != ConnectionStatus.CONNECTED or self._client is None:
            return False

        try:
            # Try to read holding register 0 as health check
            result = await self._client.read_holding_registers(
                address=0, count=1, slave=self.unit_id
            )

            if result.isError():
                self.logger.warning(f"Health check failed: {result}")
                return False

            return True

        except Exception as e:
            self.logger.error(f"Health check failed: {e}")
            return False

    def _parse_address(self, address: str) -> tuple[str, int]:
        """
        Parse Modbus address format

        Args:
            address: Address string (e.g., "holding:100", "coil:10", "50")

        Returns:
            Tuple of (register_type, address)

        Raises:
            ValueError: If address format is invalid
        """
        if ":" in address:
            parts = address.split(":", 1)
            register_type = parts[0].lower()
            reg_address = int(parts[1])
        else:
            # Default to holding register
            register_type = "holding"
            reg_address = int(address)

        # Validate register type
        valid_types = ["coil", "discrete", "holding", "input"]
        if register_type not in valid_types:
            raise ValueError(
                f"Invalid register type: {register_type}. "
                f"Must be one of: {', '.join(valid_types)}"
            )

        return register_type, reg_address

    async def _read_register(self, register_type: str, address: int) -> dict[str, Any]:
        """
        Read a single register based on type

        Args:
            register_type: Type of register (coil, discrete, holding, input)
            address: Register address

        Returns:
            Dictionary with 'value' and 'success' keys
        """
        if self._client is None:
            raise DeviceOfflineError("Client not initialized")

        try:
            if register_type == "coil":
                # Read coil (function code 01)
                result = await self._client.read_coils(address=address, count=1, slave=self.unit_id)
                if result.isError():
                    return {"success": False, "value": None, "error": str(result)}
                return {"success": True, "value": bool(result.bits[0])}

            elif register_type == "discrete":
                # Read discrete input (function code 02)
                result = await self._client.read_discrete_inputs(
                    address=address, count=1, slave=self.unit_id
                )
                if result.isError():
                    return {"success": False, "value": None, "error": str(result)}
                return {"success": True, "value": bool(result.bits[0])}

            elif register_type == "holding":
                # Read holding register (function code 03)
                result = await self._client.read_holding_registers(
                    address=address, count=1, slave=self.unit_id
                )
                if result.isError():
                    return {"success": False, "value": None, "error": str(result)}
                return {"success": True, "value": int(result.registers[0])}

            elif register_type == "input":
                # Read input register (function code 04)
                result = await self._client.read_input_registers(
                    address=address, count=1, slave=self.unit_id
                )
                if result.isError():
                    return {"success": False, "value": None, "error": str(result)}
                return {"success": True, "value": int(result.registers[0])}

            else:
                raise ValueError(f"Unknown register type: {register_type}")

        except Exception as e:
            return {"success": False, "value": None, "error": str(e)}

    async def _write_register(self, register_type: str, address: int, value: Any) -> bool:
        """
        Write a single register based on type

        Args:
            register_type: Type of register (coil, holding)
            address: Register address
            value: Value to write

        Returns:
            bool: True if write successful
        """
        if self._client is None:
            raise DeviceOfflineError("Client not initialized")

        try:
            if register_type == "coil":
                # Write single coil (function code 05)
                result = await self._client.write_coil(
                    address=address, value=bool(value), slave=self.unit_id
                )
                return not result.isError()

            elif register_type == "holding":
                # Write single holding register (function code 06)
                result = await self._client.write_register(
                    address=address, value=int(value), slave=self.unit_id
                )
                return not result.isError()

            elif register_type in ["discrete", "input"]:
                # This should not be reached due to pre-check in write_single
                return False

            else:
                raise ValueError(f"Unknown register type: {register_type}")

        except Exception as e:
            self.logger.error(f"Error writing register: {e}")
            return False
